#pragma once

#include "NumberTheory.h"

template<class T>
struct RSAKeyPair
{
private:
	T n;
	T e;
	T d;
	T p;
	T q;
	T f;

public:
	RSAKeyPair() :state(NOPARAM) {}
	int getState()const { return state; }
	int setParam(const T& _p, const T& _q);
	int setKey();
	int setKey(const T& pubkey);
	int encrypt(const T& m, T& c)const;
	int decrypt(const T& c, T& m)const;
	const T& getN()const { return n; }
	const T& getE()const { return e; }
	const T& getD()const { return d; }

private:
	int state;
	int mSetKey();

public:
	static const int ENABLE = 0;
	static const int NOKEY = 1;
	static const int NOPARAM = 2;
};

template<class T>
int RSAKeyPair<T>::setParam(const T& _p, const T& _q)
{
	if (!isPrime(_p)) return -1;
	if (!isPrime(_q)) return -2;
	p = _p;
	q = _q;
	f = (_p - 1)*(_q - 1);
	n = _p*_q;
	state = NOKEY;
	return 0;
}

template<class T>
int RSAKeyPair<T>::setKey()
{
	if (state == NOPARAM) return NOPARAM;

	e = p + 1;
	while (e < f&&gcd(e, f) != 1) ++e;
	if (e >= f) return -1;
	else
		return mSetKey();
}

template<class T>
int RSAKeyPair<T>::setKey(const T& pubkey)
{
	if (state == NOPARAM) return NOPARAM;
	if (pubkey >= f || gcd(pubkey, f) != 1) return -1;

	e = pubkey;
	return mSetKey();
}

template<class T>
int RSAKeyPair<T>::mSetKey()
{
	d = invmod(e, f);
	if (d <= 0) return -2;
	
	state = ENABLE;
	return 0;
}

template<class T>
int RSAKeyPair<T>::encrypt(const T& m,T& c)const
{
	if (state != ENABLE) return -1;
	c = powmod(m, e, n);
	return 0;
}

template<class T>
int RSAKeyPair<T>::decrypt(const T& c, T& m)const
{
	if (state != ENABLE) return -1;
	m = powmod(c, d, n);
	return 0;
}